# coding=utf-8
# pip install tensorflow-hub
# pip install bert-for-tf2
"""
参考上面的项目bert_testing改动而来
https://github.com/huwenxianglyy/bert-use-demo/blob/master/demo3.py
"""


import bert
import tensorflow as tf
import numpy as np
import create_input
import tokenization
tf.enable_eager_execution()

batch_size=20
num_labels=2  # 类别数量，我的例子是个二分类
is_training=True
max_seq_length=128
iter_num=1000
lr=0.00005

base_path = "E:/___机器学习数据集/"
bert_dir = base_path + "uncased_L-2_H-128_A-2" # JM: 这里由于原来的BERT与训练模型太大，8G显存装不下，所以改用小一号版本

vocab_file=base_path + "uncased_L-2_H-128_A-2/vocab.txt"

bert_params = bert.params_from_pretrained_ckpt(bert_dir)
if max_seq_length > bert_params.max_position_embeddings: # 模型有个最大的输入长度 512
    raise ValueError("超出模型最大长度")

# 加载数据集合
with open("data/text.txt","r",encoding="utf-8") as reader:
    data=reader.read().splitlines()

texts=[]
labels=[]
for line in data:
    line=line.split("\t")
    if len(line) ==2 and int(line[1])<2:  # 这里演示一个二分类问题，但训练样本并没有认真处理过，所以去掉label大于1的。
        texts.append(line[0])
        labels.append(line[1])

tokenizer = tokenization.FullTokenizer(vocab_file=vocab_file) # token 处理器，主要作用就是 分字，将字转换成ID。vocab_file 字典文件路径
input_idsList=[]
input_masksList=[]
segment_idsList=[]
for t in texts:
    single_input_id, single_input_mask, single_segment_id=create_input.convert_single_example(max_seq_length,tokenizer,t)
    input_idsList.append(single_input_id)
    input_masksList.append(single_input_mask)
    segment_idsList.append(single_segment_id)

#  创建bert的输入
input_ids = np.asarray(input_idsList,dtype=np.int32)
input_mask = np.asarray(input_masksList,dtype=np.int32)
segment_ids = np.asarray(segment_idsList,dtype=np.int32)
labels = np.asarray(labels,dtype=np.int32)

# 创建bert模型
l_bert = bert.BertModelLayer.from_params(bert_params, name="bert")
output = l_bert([input_ids, segment_ids])  # JM: 这里就是讲要分类的句子转换成了BERT的数字，所谓的预训练完成


# JM: 后面再接一个全连接，或者跟TextCNN应该也可以
# JM: 其实仔细想想，BERT后面接入TextCNN网络是多此一举画蛇添足，因为BERT本身已经替代了TextCNN全连接前的网络，只要最后进行分类即可
# 这不就是预训练的好处吗
print("done")